import time
import random
from typing import Any

class SnowflakeIDGenerator:
    def __init__(self, worker_id: int=1, datacenter_id: int = 1, epoch: int = 1609459200000):
        # 定义位数
        self.worker_id_bits = 5
        self.datacenter_id_bits = 5
        self.sequence_bits = 12
        
        # 最大值
        self.max_worker_id = -1 ^ (-1 << self.worker_id_bits)
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_id_bits)
        self.max_sequence = -1 ^ (-1 << self.sequence_bits)
        
        # 位移量
        self.worker_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.worker_id_bits
        self.timestamp_shift = self.sequence_bits + self.worker_id_bits + self.datacenter_id_bits
        
        # 自定义时间戳（避免与系统时间冲突）
        self.epoch = epoch
        
        # 机器ID和数据中心ID
        if worker_id > self.max_worker_id or worker_id < 0:
            raise ValueError(f"worker_id must be between 0 and {self.max_worker_id}")
        self.worker_id = worker_id
        self.datacenter_id = datacenter_id
        
        # 序列号（同一毫秒内的递增序列号）
        self.sequence = 0
        self.last_timestamp = -1
    
    def _current_timestamp(self) -> int:
        """获取当前时间戳"""
        return int(time.time() * 1000) - self.epoch
    
    def _wait_for_next_timestamp(self, last_timestamp: int) -> int:
        """等待下一毫秒"""
        timestamp = self._current_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._current_timestamp()
        return timestamp
    
    def generate_id(self) -> int:
        """生成唯一 ID"""
        timestamp = self._current_timestamp()
        
        if timestamp == self.last_timestamp:
            self.sequence = (self.sequence + 1) & self.max_sequence
            if self.sequence == 0:
                timestamp = self._wait_for_next_timestamp(self.last_timestamp)
        else:
            self.sequence = 0
        
        self.last_timestamp = timestamp
        
        # 组合所有部分，生成唯一 ID
        id = ((timestamp) << self.timestamp_shift) | (self.datacenter_id << self.datacenter_id_shift) | (self.worker_id << self.worker_id_shift) | self.sequence
        return id

generator = SnowflakeIDGenerator(worker_id=1, datacenter_id=1)